import os
import torch
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import random
from tqdm import tqdm
import argparse
import warnings
from concurrent.futures import ThreadPoolExecutor
warnings.filterwarnings('ignore', category=UserWarning)

from utils.common import check_and_create_directory, list_all_files, chunk_list
from utils.auto_sam import auto_sam
from utils.vit.vit import vit_encode

def process_file(f, out_path):
    masks = torch.load(f)
    new_pt = torch.ones_like(masks[0]['segmentation']) * -1
    for idx, mask in enumerate(masks):
        seg_mask = mask['segmentation']
        new_pt[seg_mask] = idx
    name = os.path.split(f)[-1]
    new_f = os.path.join(out_path, name)
    torch.save(new_pt, new_f)

if __name__ == '__main__':
    path = '/home/zry/datasets/building/train/tmp'
    out_path = '/home/zry/datasets/building/train/tmp_bak'
    check_and_create_directory(out_path)

    fs = list_all_files(path, '.pt')
    with ThreadPoolExecutor(max_workers=2) as executor:  # Use 4 threads
        for i, f in tqdm(enumerate(fs), total=len(fs)):
            executor.submit(process_file, f, out_path)
